package gorm

import (
	"errors"
	"fmt"
	"log"

	"mcp-server/internal"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

const Version = "1.0.0"

func init() {
	internal.GlobalRegistry.RegisterServer("database", &DBConfig{})
}

type DBConfig struct {
	dbType   string
	dsn      string
	readOnly bool
}

func (c *DBConfig) ParseConfig(config map[string]any) error {
	dsn, ok := config["dsn"].(string)
	if !ok {
		return errors.New("missing dsn")
	}
	c.dsn = dsn

	dbType, ok := config["dbType"].(string)
	if !ok {
		return errors.New("missing database type")
	}
	c.dbType = dbType

	readOnly, ok := config["readOnly"].(bool)
	if !ok {
		readOnly = false
	}
	c.readOnly = readOnly
	log.Printf("DBConfig ParseConfig: %+v", config)
	return nil
}

func (c *DBConfig) NewServer(serverName string) (*server.MCPServer, error) {
	mcpServer := server.NewMCPServer(
		serverName,
		Version,
		server.WithInstructions(fmt.Sprintf("This is a %s database server", c.dbType)),
	)

	dbClient, err := NewDBClient(c.dsn, c.dbType)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize DBClient: %w", err)
	}

	// Add query tool
	// mcpServer.AddTool(
	// 	mcp.NewToolWithRawSchema("query", fmt.Sprintf("Run a read-only SQL query in database %s", c.dbType), GetQueryToolSchema()),
	// 	HandleQueryTool(dbClient),
	// )

	dbType := dbClient.db.Dialector.Name()
	mcpServer.AddTool(createReadQueryTool(dbType), ReadQueryToolHandler(dbClient))
	mcpServer.AddTool(createListTablesTool(dbType), ListTableToolHandler(dbClient))
	if !c.readOnly {
		mcpServer.AddTool(createWriteQueryTool(dbType), WriteQueryToolHandler(dbClient))
		mcpServer.AddTool(createCreateTableTool(dbType), CreateTableToolHandler(dbClient))
	}

	return mcpServer, nil
}

func createReadQueryTool(dbType string) mcp.Tool {
	return mcp.NewTool("read_query",
		mcp.WithDescription("Execute a SELECT query on the "+dbType+" database"),
		mcp.WithString("query",
			mcp.Required(),
			mcp.Description("SELECT SQL query to execute"),
		),
	)
}

func createWriteQueryTool(dbType string) mcp.Tool {
	return mcp.NewTool("write_query",
		mcp.WithDescription("Execute an INSERT, UPDATE, or DELETE query on the "+dbType+" database"),
		mcp.WithString("query",
			mcp.Required(),
			mcp.Description("SQL query to execute"),
		),
	)
}

func createCreateTableTool(dbType string) mcp.Tool {
	return mcp.NewTool("create_table",
		mcp.WithDescription("Create a new table in the "+dbType+" database"),
		mcp.WithString("schema",
			mcp.Required(),
			mcp.Description("CREATE TABLE SQL statement"),
		),
	)
}

func createListTablesTool(dbType string) mcp.Tool {
	return mcp.NewTool("list_tables",
		mcp.WithDescription("List all user tables in the database"),
		mcp.WithString("schema",
			mcp.Description("Optional schema name to filter tables"),
		),
	)
}
